import concurrent.futures
import atexit
import importlib
import os
import sys
import time
import json
import matplotlib.pyplot as plt
import numpy as np
import warnings
import shutil
import random

warnings.filterwarnings('ignore', category=FutureWarning)
#matplotlib.use('TkAgg')

from a1_planner import Planner as Planner_a1
from a2_planner import Planner as Planner_a2
from lqr_control import LQRController
from a1_simple_control import MotionController as MotionController_a1
from a2_simple_control import MotionController as MotionController_a2
from predictor import Predictor
from planner_backup import process as Planner_backup
from get_traj_true_dict import MotionController as MotionController_a11
from planner_backup import is_valid_scene_name


def exit_handler(executor):
    print("Exiting program...")
    executor.shutdown(wait=False)


class Logger:
    def __init__(self):
        self.terminal = sys.stdout
        self.log = []  # 用于存储print输出的内容

    def write(self, message):
        self.terminal.write(message)  # 输出到标准输出
        self.log.append(message)  # 记录print输出的内容

    def flush(self):
        pass  # 不执行任何操作，仅为了满足标准输出的要求

# 创建Logger的实例
logger = Logger()


def test_scenario(scenario_to_test, dir_save_img, dir_parent_1):
    try:
        random.seed(0)  # 设置随机种子
        dir_current_file = os.path.dirname(__file__)  # 'Algorithm_demo_1\planner'
        dir_parent_1 = os.path.dirname(dir_current_file)  # 'Algorithm_demo_1'

        # save_dir = os.path.join(os.path.dirname(dir_parent_1), "all_results", "path_spd_saved")
        save_dir = os.path.join(dir_current_file, "all_results", "path_spd_saved")
        os.makedirs(save_dir, exist_ok=True)

        # 在子进程中初始化模块和对象
        scenarioOrganizer = importlib.import_module("onsite-mine.dynamic_scenes.scenarioOrganizer1")
        env = importlib.import_module("onsite-mine.dynamic_scenes.env")
        lookup = importlib.import_module("onsite-mine.dynamic_scenes.lookup")
        envi = env.Env()

        so = scenarioOrganizer.ScenarioOrganizer()
        dir_inputs = os.path.abspath(os.path.join(dir_parent_1, 'inputs'))
        dir_outputs = os.path.abspath(os.path.join(dir_parent_1, 'outputs'))
        so.load(dir_inputs, dir_outputs)

        scene_name = scenario_to_test['data']['scene_name']
        print(f"###log### <scene-{scene_name}>\n")
        formatted_so_config = json.dumps(so.config, indent=4, ensure_ascii=False)
        print(f"###log### <测试参数>\n{formatted_so_config}\n")
        collision_lookup = lookup.CollisionLookup()
        observation, traj, client = envi.make(scenario=scenario_to_test, save_img_path=dir_save_img,
                                              kinetics_mode='simple')
        # traj_true_dict = {}
        path_planned_1 = None
        replan = False
        predictor = Predictor(time_horizon=5.0)
        lat_controller = MotionController_a1()
        lat_controller1 = MotionController_a11()
        if observation['test_setting']['scenario_type'] == 'intersection':
            ########### 算法 部分1:初始化 ###########
            # 交叉路口单车通行任务:构建整个场景的预测器,预测所有的车辆;构建规划器\控制器,控制ego车
            planner = Planner_a1(observation)
            save_file = os.path.join(save_dir, f"{scene_name}.npz")
            if os.path.exists(save_file):
                data = np.load(save_file)
                path_planned = data['path_planned'].tolist()
                # print(f"加载文件{save_file}")
            else:
                path_planned, spd_planned = planner.process(observation)
                np.savez(save_file, path_planned=path_planned, spd_planned=spd_planned)
            # step_sum = int(observation["test_setting"]['max_t'] / observation["test_setting"]['dt'])
            ########### 算法 部分1:初始化 ###########

            # 逐帧进行仿真，触发仿真停止条件时结束
            # 当测试还未进行完毕,即观察值中test_setting['end']还是-1的时候
            while observation['test_setting']['end'] == -1:
                traj_future, _ = predictor.predict(observation, traj, predictor_type="CACV_PREDICTOR")
                if scene_name not in ['dapai_intersection_2_2_4','dapai_intersection_1_2_17',
                                      'dapai_intersection_1_3_13']:
                    _, traj_true_dict = lat_controller1.process(observation['vehicle_info'], path_planned, traj_future,
                                                                observation,
                                                                traj)
                    ########### 算法 部分2:执行 ###########
                    time_simu = round(observation['test_setting']['t'], 1)
                    if scene_name not in ['jiangtong_intersection_1_3_2','jiangtong_intersection_6_3_1', 'jiangtong_intersection_6_3_6', 'jiangtong_intersection_1_2_6',
                                          'jiangtong_intersection_1_3_10', 'jiangtong_intersection_5_3_2',
                                          'jiangtong_intersection_1_2_5'] :
                        if time_simu % 1 == 0 and path_planned_1 is None:   # 不动态规划
                            # print(f"time_simu={time_simu:.3f}s")
                            path_planned_1, _ = Planner_backup(scenario_to_test, observation, traj_true_dict, traj)
                    else:
                        if time_simu % 1 == 0:  # 动态规划
                            # print(f"time_simu={time_simu:.3f}s")
                            path_planned_1, _ = Planner_backup(scenario_to_test, observation, traj_true_dict, traj)
                    # if traj_true_dict is not None and len(traj_true_dict) == 1 and path_planned_1 is None:
                    #     path_planned_1, _ = Planner_backup(scenario_to_test, observation, traj_true_dict, traj)
                    # replan = False  # 动态规划中适用,没有规划好，慢点跑
                    if path_planned_1 is not None:
                        path_planned = path_planned_1
                        replan = True
                action, _ = lat_controller.process(observation['vehicle_info'], path_planned, traj_future, observation,
                                                traj,replan=replan)
                action = action + (1,)  # 添加档位信息，默认为前进档位
                ########### 算法 部分2:执行 ###########
                observation = envi.step(action, traj_future, observation, traj)  # 根据车辆的action,更新场景,并返回新的观测值.
        else:
            Reverse = False
            simple = False
            back_move = False
            switch = False
            planner = Planner_a2(observation)

            save_file = os.path.join(save_dir, f"{scene_name}.npz")
            if os.path.exists(save_file):
                data = np.load(save_file)
                path_planned = data['path_planned'].tolist()
                spd_planned = data['spd_planned'].tolist()
                # print(f"加载文件{save_file}")
            else:
                path_planned, spd_planned = planner.process(collision_lookup, observation,traj=traj)
                np.savez(save_file, path_planned=path_planned, spd_planned=spd_planned)

            task = scene_name[scene_name.find("_") + 1:scene_name.find("_", scene_name.find("_") + 1)]
            ind_list = []
            while observation['test_setting']['end'] == -1:
                if task == "loading":
                    if Reverse == False:
                        controller = LQRController()
                        action, Reverse, ind_list = controller.process(observation['vehicle_info'], path_planned,
                                                                       spd_planned,
                                                                       "simple", Reverse=Reverse, simple=simple,
                                                                       ind_list=ind_list)
                    else:
                        controller = MotionController_a2(20)
                        action, ind, Reverse, arrive,_ = controller.process(observation['vehicle_info'],
                                                                          path_planned,
                                                                          spd_planned, observation,
                                                                          Reverse=Reverse,
                                                                          traj=traj,
                                                                          scenario_to_test=scenario_to_test,switch=switch)
                        if arrive:
                            Reverse = False
                            simple = True
                            print("切换控制方法为LQR")
                elif task == "unloading":
                    if traj:
                        for car in traj:
                            back_move = any(value.get('v_mps', 0) != 0 for key, value in traj[car].items() if isinstance(value, dict))

                    if back_move and not is_valid_scene_name(scene_name):   # 如果背景车在动
                        traj_future, _ = predictor.predict(observation, traj, predictor_type="CACV_PREDICTOR")
                        # 冗余
                        _, traj_true_dict = lat_controller1.process(observation['vehicle_info'], path_planned, traj_future, observation,
                                                    traj)

                        time_simu = round(observation['test_setting']['t'], 1)
                        if time_simu % 1 == 0 :
                            # print(f"time_simu={time_simu:.3f}s")
                            path_planned_1, spd_planned_1 = Planner_backup(scenario_to_test, observation, traj_true_dict, traj,back_move=back_move)

                        # if traj_true_dict is not None and len(traj_true_dict) == 1 and path_planned_1 is None:
                        #     path_planned_1, _ = Planner_backup(scenario_to_test, observation, traj_true_dict, traj)
                        if path_planned_1 is not None:
                            path_planned = path_planned_1
                            spd_planned = spd_planned_1


                    controller = MotionController_a2(10)
                    action, _, _, _,_ = controller.process(observation['vehicle_info'], path_planned,
                                                         spd_planned, observation, Reverse=False, traj=traj,
                                                         scenario_to_test=scenario_to_test,switch=False)
                observation = envi.step(action, 100, observation, 100)

    except Exception as e:
        print(repr(e))
        if client is not None:
            client.close_sockets()
    finally:
        so.add_result(scenario_to_test, observation['test_setting']['end'])
        plt.close()


def main():
    dir_current_file = os.path.dirname(__file__)
    dir_parent_1 = os.path.dirname(dir_current_file)
    dir_save_img = os.path.abspath(os.path.join(dir_parent_1, 'onsite_images_saved'))

    file_path = os.path.join(os.path.dirname(dir_parent_1), "all_results", 'print_log.txt')
    os.makedirs(os.path.dirname(file_path), exist_ok=True)

    tic = time.time()

    # 加载场景
    scenarioOrganizer = importlib.import_module("onsite-mine.dynamic_scenes.scenarioOrganizer1")
    so = scenarioOrganizer.ScenarioOrganizer()
    dir_inputs = os.path.abspath(os.path.join(dir_parent_1, 'inputs'))
    dir_outputs = os.path.abspath(os.path.join(dir_parent_1, 'outputs'))
    so.load(dir_inputs, dir_outputs)

    # 准备测试场景列表
    scenario_list = []
    while True:
        scenario_to_test = so.next()
        if scenario_to_test is None:
            break
        scenario_list.append(scenario_to_test)

    # for scenario_to_test in scenario_list:
    #     test_scenario(scenario_to_test, dir_save_img, dir_parent_1)

    executor = concurrent.futures.ProcessPoolExecutor(max_workers=4)
    atexit.register(exit_handler, executor)
    with executor as executor_inner:
        futures = []
        for scenario_to_test in scenario_list:
            futures.append(executor_inner.submit(test_scenario, scenario_to_test, dir_save_img, dir_parent_1))
        for future in concurrent.futures.as_completed(futures):
            pass
    toc = time.time()
    print(f"###log### 总用时: {int((toc - tic) // 3600)}小时 {int(((toc - tic) % 3600) // 60)}分钟 {(toc - tic) % 60:.2f}秒\n")
    if not os.path.exists(file_path):
        with open(file_path, 'w') as f:
            pass  # 什么都不做，只是创建文件

    # 打开文件并读取内容
    with open(file_path, 'r') as f:
        existing_content = f.read()

    # 将 "output.log" 文件的内容追加到 "all_result.txt" 文件中
    with open(file_path, 'a') as f:
        # 添加长横线和换行进行区分
        f.write("\n")
        f.write("-" * 200)
        f.write("\n")
        for message in logger.log:
            if not message.startswith("###log### <测试参数>") and not message.startswith("{"):
                f.write(message)

    out_dir = os.path.join(os.path.dirname(os.path.dirname(file_path)), "outputs")
    os.makedirs(out_dir, exist_ok=True)
    for filename in os.listdir(dir_outputs):
        if filename.endswith('.csv'):
            src_file = os.path.join(dir_outputs, filename)
            dst_file = os.path.join(out_dir, filename)
            shutil.copy2(src_file, dst_file)


if __name__ == "__main__":
    random.seed(0)  # 设置随机种子
    sys.stdout = logger
    main()
